import marimo

__generated_with = "0.15.3"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    return (mo,)


@app.cell
def _(mo):
    mo.md(
        r"""
    # Transformer简单解释：从QKV到多头注意力机制

    如果你曾经尝试阅读具有里程碑意义的"Attention Is All You Need"论文，你一定知道那种感觉。你跟着图表走，理解了高层概念，但随后就撞墙了。对我来说，那堵墙就是Transformer的核心机制：**Query**、**Key**和**Value**到底代表什么？我多次阅读论文，但没有对这些组件的扎实直觉，架构感觉就像一系列复杂的黑盒子。今天，我将穿透数学，向你展示Transformer魔力背后简单而优雅的想法。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ### 以前深度学习模型的挑战

    Transformer的最大突破不仅在于它们的性能，还在于它们完全改变了模型"看"序列的方式。它们引入了一个强大的概念叫做**注意力机制**。

    在我们深入注意力机制的细节之前，让我们思考一下它解决的问题。**如果传统的深度学习模型读取一个长句子，它会按顺序逐词处理。**

    _你认为当该模型试图理解句子开头的一个词与远在句子末尾的词之间的关系时，主要挑战是什么？_
    """
    )
    return


@app.cell
def _(mo):
    mo.image(
        src="https://miro.medium.com/v2/resize:fit:700/1*Va6y0eOgruBK_k3mfm5kYA.png",
        alt="传统模型处理序列的挑战"
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    像RNN和LSTM这样的旧模型面临的一个大挑战是它们的顺序逐词处理。这造成了两个主要问题：

    1. **速度：** 一次处理一个词非常慢，特别是对于长文档。
    2. **长期记忆：** 随着句子变长，模型必须记住开头发生的事情。来自早期词汇的信息可能会在处理新词时"淡化"，使得很难将它们与后面的词连接起来。
    """
    )
    return


@app.cell
def _(mo):
    mo.image(
        src="https://miro.medium.com/v2/resize:fit:700/1*dSFfG1Gs-Ep-kO2tvhzusA.png",
        alt="简单RNN架构（显示当前状态与前一状态的依赖关系）"
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ### Transformer如何实现并行处理

    **注意力机制**就是解决这个问题的方法。注意力机制不采用逐词方法，而是让模型一次性查看句子中的所有词，并决定哪些词对理解特定词最重要。

    > 让我们用一个例子：The animal didn't cross the street because **it** was too tired.（动物没有过马路，因为**它**太累了。）

    当模型处理词"**it**"时，你认为它需要对哪个词给予最多的"注意力"才能理解"it"指的是什么？
    """
    )
    return


@app.cell
def _(mo):
    mo.image(
        src="https://miro.medium.com/v2/resize:fit:700/0*OLBkC21ZcCYDQ292.gif",
        alt="注意力机制示例"
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    这就是Transformer的**自注意力**机制的核心思想。

    模型不只是按顺序读取词汇。当它处理像"**it**"这样的词时，它会创建与句子中每个其他词的连接，并给它们一个"注意力分数"来找出哪些最相关。对于"**it**"，词"animal"得到很高的分数。词"street"得到低分数。

    计算这些分数的过程涉及每个词的三个关键概念：**Query**、**Key**和**Value**。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ### Q、K和V背后的直觉

    让我们用一个简单的类比。把它想象成一个图书馆。你的**Query**是你在计算机中输入的搜索词。图书馆的卡片目录有一堆**Keys**（如书名或主题）。当一个**Key**匹配你的**Query**时，你得到一个**Value**（实际的书）。

    在我们的句子中，每个词都充当自己的**Query**、**Key**和**Value**。

    当模型处理词"**it**"时，它使用其**Query**来"询问"每个其他词的**Key**。

    你认为"**it**"的**Query**在**Key**中寻找什么？

    模型被训练学习像"it"这样的代词的**Query**在语义上类似于像"animal"这样的名词的**Key**。它在寻找一个说"嘿，我是一个单数、非人类名词，可能是那个动作的主语"的键。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    现在，一旦模型发现"animal"的注意力分数非常高，下一个逻辑步骤是什么？词"animal"需要向词"it"提供什么信息，以便"it"能够被完全理解？这就是**Value**发挥作用的地方。

    "animal"的**Value**包含关于该词的所有学习到的描述性信息。

    在深度学习的世界中，这种"描述性信息"不是作为像"非人类"或"名词"这样的词汇列表存储的。相反，它是一个称为**词嵌入**或**向量**的数字列表。这个向量是词汇含义的数值表示，捕获其所有上下文、属性和与其他词的关系。

    所以，回顾词"**it**"的整个**自注意力**过程：

    1. "**it**"的**Query**询问："什么是可能是'was too tired'主语的单数、非人类名词？"
    2. 它将该查询与所有其他词的**Keys**进行比较。"animal"的**Key**有很好的匹配。
    3. 模型使用该匹配给"animal"一个高**注意力分数**。
    4. 最后，它取"animal"的**Value**（丰富的数值含义）并将其发送回"**it**"，允许"**it**"在其适当的上下文中被理解。

    这整个过程对句子中的_每个_词都进行，**同时进行**！这就是为什么Transformer如此快速且如此擅长理解长距离关系。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    现在，Transformer架构的一个关键部分是它不只做一次。它使用称为**多头注意力**的东西并行地多次执行。你认为为什么模型同时从几个不同的"视角"看同一个句子会有帮助？_在跳到下一段之前先思考一下。_

    多头注意力完全是关于捕获词汇之间关系的不同"维度"或视角。

    这样想：

    - 一个**注意力头**可能学会专注于**语法关系** — 例如，找出词"tired"是"animal"的描述。
    - 另一个**注意力头**可能专注于识别**共指** — 比如找出"it"指的是"animal"。
    - 第三个头可能完全专注于另一种关系，比如句子中主语和宾语之间的关系。

    通过同时运行所有这些"头"然后组合它们的输出，模型获得了比单一注意力计算更丰富、更完整的句子理解。这就像有一个专家团队，每个人都在看同一个问题的不同方面。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ### 模型如何解决词序问题

    因为注意力机制一次性查看句子中的所有词，它完全失去了自然顺序。没有关于词位置的任何信息，模型会将"_The dog bit the man_"（狗咬了人）视为与"_The man bit the dog_"（人咬了狗）相同，即使含义完全不同。

    为了解决这个问题，Transformer添加了一个称为**位置编码**的特殊组件。

    所以，总结我们到目前为止涵盖的内容是原始论文中给出的**编码器块**层：

    - 输入词被转换为数值表示（**嵌入**）。
    - **位置编码**被添加到这些嵌入中，为模型提供关于词序的信息。
    - 这些丰富的词向量然后通过**多头注意力**层，该层计算所有词之间的关系。
    - 最后，输出通过一个简单的**前馈网络**来处理新的、上下文丰富的信息。
    """
    )
    return


@app.cell
def _(mo):
    mo.image(
        src="https://miro.medium.com/v2/resize:fit:700/1*p84DlJFOXInmMLrvGPRJ5A.png",
        alt="原始论文中的编码器架构"
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ### 解码器

    你认为解码器面临的主要挑战是什么，而编码器不必处理？它的注意力机制如何需要稍微不同来处理这个挑战？

    解码器的工作是接受编码器的输出并生成新序列，如翻译或摘要。

    解码器必须逐词处理，因为它正在生成新序列。当它试图预测下一个词时，看到"答案"（句子的其余部分）是没有意义的。

    这导致解码器中稍微不同类型的注意力，称为**掩码自注意力**。它的工作原理就像我们之前讨论的自注意力，但它有一个"掩码"来隐藏所有未来的词。这迫使解码器只关注它已经生成的词。

    但是等等，解码器也需要了解来自编码器的_原始_句子，对吧？这就是第二种注意力类型的用武之地：**交叉注意力**。
    """
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    这样想：

    - **编码器**读取整个原始句子并创建对它的深入、上下文丰富的理解。这就像学生阅读和理解一篇文章。
    - **解码器**然后开始写那篇文章的摘要。
    - 它使用**掩码自注意力**来阅读它到目前为止写的内容，以确保摘要有意义。
    - 它使用**交叉注意力**来回顾原始文章（编码器的输出），以确保其摘要准确。

    所以，解码器有更复杂的工作，因为它必须平衡自己的输出与从编码器获得的信息。

    最后两个步骤将这些数字转换为实际词汇：

    1. **线性层：** 解码器的最终输出被馈送到一个简单的全连接神经网络层。这一层的工作是将数字向量扩展为更大的向量。这个新的、更大的向量与模型的整个词汇表（它知道的所有词）大小相同。这个向量中的每个数字对应一个特定的词。
    2. **Softmax函数：** 这是神奇的部分。Softmax函数取线性层的数字并将它们转换为概率。对应于词"cat"的数字可能变成90%的概率，而"dog"可能变成5%，等等。具有最高概率的词被选为输出序列中的下一个词。

    模型重复这个过程 — 取其自己新生成的词并将其用作下一步的输入 — 直到它生成一个表示句子结束的特殊标记。
    """
    )
    return


@app.cell
def _(mo):
    mo.image(
        src="https://miro.medium.com/v2/resize:fit:700/1*XvumALYTaB9hzKvkCtjXag.png",
        alt="Transformer架构"
    )
    return


@app.cell
def _(mo):
    mo.md(
        r"""
    ## 总结

    让我们快速回顾整个过程：

    1. **输入和编码**：原始句子被转换为数字（**嵌入**），然后添加**位置编码**以保持词序。
    2. **编码器**：编码器使用**自注意力**一次性理解输入句子中所有词之间的关系。
    3. **解码器**：解码器然后开始逐词生成新句子。它使用**掩码自注意力**来查看它已经生成的词，使用**交叉注意力**来回顾来自编码器的原始句子。
    4. **最终输出**：解码器的输出通过**线性层**和**softmax函数**处理，从词汇表中选择最佳的下一个词。

    这就是Transformer的核心工作原理！通过这种优雅的设计，它能够并行处理序列，捕获长距离依赖关系，并生成高质量的输出。
    """
    )
    return


if __name__ == "__main__":
    app.run()
